from tkinter import *
from tkinter.filedialog import askdirectory, askopenfilename
from tkinter.messagebox import showinfo, messagebox
import cv2
import numpy as np
from PIL import Image, ImageTk
from tkinter import ttk
import pygame
import time
import os
import threading
from datetime import datetime

# 初始化pygame音频
pygame.mixer.init(frequency=16000, size=-16, channels=2, buffer=4096)

# 加载级联分类器
detector = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_default.xml')
mask_detector = cv2.CascadeClassifier('xml/cascade.xml')

class MaskDetectionGUI:
    def __init__(self):
        self.camera = None
        self.is_running = False
        self.current_frame = None
        self.detection_history = []
        
        self.root = Tk()
        self.root.title('口罩检测系统')
        self.root.geometry('%dx%d' % (1000, 700))
        self.root.resizable(False, False)
        
        # 创建必要的目录
        self.create_directories()
        
        # 加载提示音
        self.load_alert_sound()
        
        self.create_widgets()
        self.create_menu()
        
    def create_directories(self):
        """创建必要的目录"""
        os.makedirs('haarcascades', exist_ok=True)
        os.makedirs('xml', exist_ok=True)
        os.makedirs('screenshots', exist_ok=True)
        os.makedirs('logs', exist_ok=True)
        
    def load_alert_sound(self):
        """加载提示音"""
        try:
            # 创建一个简单的提示音
            duration = 0.5
            sample_rate = 16000
            frequency = 800
            
            t = np.linspace(0, duration, int(sample_rate * duration), False)
            wave = np.sin(2 * np.pi * frequency * t) * 0.3
            
            # 转换为16位PCM
            audio = (wave * 32767).astype(np.int16)
            stereo_audio = np.column_stack([audio, audio])
            
            # 保存为临时文件
            import tempfile
            self.temp_sound = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
            import wave as wav
            with wav.open(self.temp_sound.name, 'wb') as wav_file:
                wav_file.setnchannels(2)
                wav_file.setsampwidth(2)
                wav_file.setframerate(sample_rate)
                wav_file.writeframes(stereo_audio.tobytes())
                
            self.alert_sound = self.temp_sound.name
            
        except Exception as e:
            print(f"无法加载提示音: {e}")
            self.alert_sound = None
            
    def create_widgets(self):
        """创建主界面组件"""
        # 主框架
        main_frame = Frame(self.root)
        main_frame.pack(fill=BOTH, expand=True, padx=10, pady=10)
        
        # 标题
        title_label = Label(main_frame, text='口罩实时检测系统', 
                           font=('微软雅黑', 24, 'bold'), fg='#2c3e50')
        title_label.pack(pady=10)
        
        # 视频显示区域
        video_frame = Frame(main_frame, bg='black', relief=SOLID, bd=2)
        video_frame.pack(fill=BOTH, expand=True, pady=10)
        
        self.video_label = Label(video_frame, bg='black')
        self.video_label.pack(fill=BOTH, expand=True)
        
        # 控制按钮区域
        control_frame = Frame(main_frame)
        control_frame.pack(fill=X, pady=10)
        
        # 按钮样式
        button_style = {
            'font': ('微软雅黑', 12, 'bold'),
            'width': 15,
            'height': 2,
            'relief': RAISED,
            'bd': 2
        }
        
        # 创建按钮
        self.start_btn = Button(control_frame, text='开始检测', 
                               command=self.start_detection, 
                               bg='#27ae60', fg='white', **button_style)
        self.start_btn.pack(side=LEFT, padx=5)
        
        self.stop_btn = Button(control_frame, text='停止检测', 
                              command=self.stop_detection, 
                              bg='#e74c3c', fg='white', **button_style,
                              state=DISABLED)
        self.stop_btn.pack(side=LEFT, padx=5)
        
        self.screenshot_btn = Button(control_frame, text='保存截图', 
                                    command=self.save_screenshot, 
                                    bg='#3498db', fg='white', **button_style)
        self.screenshot_btn.pack(side=LEFT, padx=5)
        
        self.file_btn = Button(control_frame, text='选择文件', 
                              command=self.select_file, 
                              bg='#f39c12', fg='white', **button_style)
        self.file_btn.pack(side=LEFT, padx=5)
        
        # 状态显示区域
        status_frame = Frame(main_frame, relief=SOLID, bd=1)
        status_frame.pack(fill=X, pady=5)
        
        self.status_label = Label(status_frame, text='就绪', 
                                 font=('微软雅黑', 10), fg='#7f8c8d')
        self.status_label.pack(side=LEFT, padx=10)
        
        # 检测统计
        self.stats_label = Label(status_frame, text='检测统计: 0人', 
                                font=('微软雅黑', 10), fg='#2c3e50')
        self.stats_label.pack(side=RIGHT, padx=10)
        
    def create_menu(self):
        """创建菜单栏"""
        menubar = Menu(self.root)
        self.root.config(menu=menubar)
        
        # 文件菜单
        file_menu = Menu(menubar, tearoff=0)
        menubar.add_cascade(label='文件', menu=file_menu)
        file_menu.add_command(label='选择图片/视频', command=self.select_file)
        file_menu.add_separator()
        file_menu.add_command(label='保存截图', command=self.save_screenshot)
        file_menu.add_separator()
        file_menu.add_command(label='退出', command=self.on_closing)
        
        # 帮助菜单
        help_menu = Menu(menubar, tearoff=0)
        menubar.add_cascade(label='帮助', menu=help_menu)
        help_menu.add_command(label='使用说明', command=self.show_help)
        help_menu.add_command(label='关于', command=self.show_about)
        
    def start_detection(self):
        """开始实时检测"""
        try:
            self.camera = cv2.VideoCapture(0)
            if not self.camera.isOpened():
                messagebox.showerror('错误', '无法打开摄像头')
                return
                
            self.is_running = True
            self.start_btn.config(state=DISABLED)
            self.stop_btn.config(state=NORMAL)
            self.status_label.config(text='检测中...')
            
            # 启动检测线程
            self.detection_thread = threading.Thread(target=self.detection_loop)
            self.detection_thread.daemon = True
            self.detection_thread.start()
            
        except Exception as e:
            messagebox.showerror('错误', f'启动检测失败: {str(e)}')
            
    def stop_detection(self):
        """停止检测"""
        self.is_running = False
        if self.camera:
            self.camera.release()
            self.camera = None
            
        self.start_btn.config(state=NORMAL)
        self.stop_btn.config(state=DISABLED)
        self.status_label.config(text='已停止')
        
    def detection_loop(self):
        """检测主循环"""
        while self.is_running and self.camera:
            ret, frame = self.camera.read()
            if not ret:
                break
                
            self.current_frame = frame.copy()
            processed_frame = self.process_frame(frame)
            
            # 更新显示
            self.update_video_display(processed_frame)
            
        cv2.destroyAllWindows()
        
    def process_frame(self, frame):
        """处理单帧图像"""
        # 转换为灰度图
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # 人脸检测
        faces = detector.detectMultiScale(gray, 1.1, 4)
        
        mask_count = 0
        no_mask_count = 0
        
        for (x, y, w, h) in faces:
            # 检测口罩
            roi_gray = gray[y:y+h, x:x+w]
            masks = mask_detector.detectMultiScale(roi_gray, 1.1, 4)
            
            if len(masks) > 0:
                # 佩戴口罩
                color = (0, 255, 0)  # 绿色
                label = "Mask"
                mask_count += 1
            else:
                # 未佩戴口罩
                color = (0, 0, 255)  # 红色
                label = "No Mask"
                no_mask_count +=
                
                # 播放提示音
                if self.alert_sound:
                    threading.Thread(target=self.play_alert).start()
                    
            # 绘制矩形框
            cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
            
            # 添加标签
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(frame, label, (x, y-10), font, 0.9, color, 2)
            
        # 更新统计
        total_faces = len(faces)
        self.update_stats(total_faces, mask_count, no_mask_count)
        
        # 记录检测历史
        if total_faces > 0:
            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            self.detection_history.append({
                'time': timestamp,
                'total': total_faces,
                'with_mask': mask_count,
                'without_mask': no_mask_count
            })
            
        return frame
        
    def play_alert(self):
        """播放提示音"""
        try:
            pygame.mixer.music.load(self.alert_sound)
            pygame.mixer.music.play()
        except:
            pass
            
    def update_video_display(self, frame):
        """更新视频显示"""
        # 转换颜色空间
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 调整大小以适应显示区域
        img = Image.fromarray(frame)
        img = img.resize((960, 540), Image.Resampling.LANCZOS)
        
        # 转换为PhotoImage
        imgtk = ImageTk.PhotoImage(image=img)
        
        # 更新标签
        self.video_label.imgtk = imgtk
        self.video_label.configure(image=imgtk)
        
    def update_stats(self, total, mask, no_mask):
        """更新统计信息"""
        self.stats_label.config(
            text=f'检测统计: {total}人 | 佩戴口罩: {mask}人 | 未佩戴口罩: {no_mask}人'
        )
        
    def select_file(self):
        """选择图片或视频文件"""
        file_path = askopenfilename(
            title='选择图片或视频文件',
            filetypes=[
                ('图片文件', '*.jpg *.jpeg *.png *.bmp'),
                ('视频文件', '*.mp4 *.avi *.mov'),
                ('所有文件', '*.*')
            ]
        )
        
        if file_path:
            self.process_file(file_path)
            
    def process_file(self, file_path):
        """处理选择的文件"""
        try:
            if file_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                # 处理图片
                frame = cv2.imread(file_path)
                if frame is not None:
                    processed = self.process_frame(frame)
                    self.current_frame = processed
                    self.update_video_display(processed)
                    self.status_label.config(text=f'已处理图片: {os.path.basename(file_path)}')
                    
            elif file_path.lower().endswith(('.mp4', '.avi', '.mov')):
                # 处理视频
                cap = cv2.VideoCapture(file_path)
                if cap.isOpened():
                    ret, frame = cap.read()
                    if ret:
                        processed = self.process_frame(frame)
                        self.current_frame = processed
                        self.update_video_display(processed)
                        self.status_label.config(text=f'已加载视频: {os.path.basename(file_path)}')
                    cap.release()
                    
        except Exception as e:
            messagebox.showerror('错误', f'处理文件失败: {str(e)}')
            
    def save_screenshot(self):
        """保存截图"""
        if self.current_frame is not None:
            try:
                timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
                filename = f'screenshot_{timestamp}.jpg'
                filepath = os.path.join('screenshots', filename)
                
                cv2.imwrite(filepath, self.current_frame)
                
                # 保存检测日志
                if self.detection_history:
                    log_filename = f'log_{timestamp}.txt'
                    log_filepath = os.path.join('logs', log_filename)
                    with open(log_filepath, 'w', encoding='utf-8') as f:
                        f.write('口罩检测日志\n')
                        f.write('=' * 50 + '\n')
                        for entry in self.detection_history[-10:]:  # 只保存最近10条
                            f.write(f"时间: {entry['time']}\n")
                            f.write(f"总人数: {entry['total']} | "
                                   f"佩戴口罩: {entry['with_mask']} | "
                                   f"未佩戴口罩: {entry['without_mask']}\n")
                            f.write('-' * 30 + '\n')
                            
                messagebox.showinfo('成功', f'截图已保存: {filename}')
                
            except Exception as e:
                messagebox.showerror('错误', f'保存截图失败: {str(e)}')
        else:
            messagebox.showwarning('警告', '没有可保存的画面')
            
    def show_help(self):
        """显示帮助信息"""
        help_text = """
使用说明：

1. 开始检测：点击"开始检测"按钮启动摄像头实时检测
2. 停止检测：点击"停止检测"按钮停止检测
3. 保存截图：点击"保存截图"按钮保存当前画面
4. 选择文件：点击"选择文件"按钮处理本地图片或视频
5. 检测结果：绿色框表示已佩戴口罩，红色框表示未佩戴口罩

注意事项：
- 确保摄像头权限已开启
- 检测模型文件需要放置在对应目录
- 截图和日志会自动保存到对应文件夹
        """
        messagebox.showinfo('使用说明', help_text)
        
    def show_about(self):
        """显示关于信息"""
        about_text = """
口罩实时检测系统 v1.0

基于OpenCV和Tkinter开发的口罩检测系统
支持实时摄像头检测和本地文件处理

技术栈：
- Python 3.10
- OpenCV-Python 4.8
- Tkinter 8.6
- TensorFlow 2.13
- pygame 2.5
        """
        messagebox.showinfo('关于', about_text)
        
    def on_closing(self):
        """关闭窗口时的处理"""
        if self.is_running:
            self.stop_detection()
        self.root.destroy()
        
    def run(self):
        """运行程序"""
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        self.root.mainloop()


if __name__ == '__main__':
    # 检查必要的文件
    if not os.path.exists('haarcascades/haarcascade_frontalface_default.xml'):
        print("警告：人脸检测模型文件不存在，请下载并放置到 haarcascades 目录")
        
    if not os.path.exists('xml/cascade.xml'):
        print("警告：口罩检测模型文件不存在，请下载并放置到 xml 目录")
        
    # 启动程序
    app = MaskDetectionGUI()
    app.run()
